#!/usr/bin/python2
# -*- coding: utf-8 -*-
# Name:     mymoldb.py
# Author:   xiooli <xioooli[at]yahoo.com.cn>
# Version:  111202

import pycurl, StringIO, urllib, json, tempfile

webapi_url = "http://210.75.237.226/db/webapi"
login_url = "http://210.75.237.226/db/user/login"
options = []
username = ''
password = ''
proxy = ''

class Query():
    def __init__(self, webapi_url = webapi_url, login_url = login_url,
            options = options, username = username,
            password = password, proxy = proxy):

        self.host = webapi_url.replace('//', '$$').split('/')[0].replace('$$', '//')
        self.webapi_url = webapi_url
        self.proxy = proxy
        self.login_url = login_url

        self.crl_post = pycurl.Curl()
        self.crl_get = pycurl.Curl()
        self.fp = StringIO.StringIO()
        #self.cookiefile = '/tmp/a.txt'
        self.cookiefile = tempfile.mkstemp(text=True)[1]
        if options:
            self.options = options
        else:
            self.options = [
                    (pycurl.VERBOSE, 0),
                    (pycurl.FOLLOWLOCATION, 1),
                    (pycurl.MAXREDIRS, 5),
                    #(pycurl.AUTOREFERER,1),
                    (pycurl.CONNECTTIMEOUT, 60),
                    (pycurl.TIMEOUT, 120),
                    (pycurl.COOKIEJAR, self.cookiefile),
                    (pycurl.COOKIEFILE, self.cookiefile),
                    (pycurl.HTTPPROXYTUNNEL, 1),
                    (self.crl_post.WRITEFUNCTION, self.fp.write),
                    #(pycurl.NOQtCore.SIGNAL, 1),
                    (pycurl.USERAGENT, "SmartMass Confirmor"),
                    (pycurl.HTTPHEADER, ["Accept: */*",
                        "Accept-Charset: ISO-8859-1,utf-8;q=0.7,*;q=0.7",
                        "Connection: keep-alive",
                        "Expect:",
                        "Keep-Alive: 300"])
                    ]

        if self.proxy:
            self.options.append((pycurl.PROXY, self.proxy))

        for opt in self.options:
            self.crl_post.setopt(opt[0], opt[1])
            self.crl_get.setopt(opt[0], opt[1])

        self.crl_post.setopt(pycurl.URL, self.webapi_url)

        self.login(username, password)

    def login(self, username = '', password = ''):
        session_crl = pycurl.Curl()

        for opt in self.options:
            session_crl.setopt(opt[0], opt[1])

        if self.login_url and username and password:
            session_crl.setopt(pycurl.URL, self.login_url)
            session_crl.setopt(session_crl.POSTFIELDS,
                    urllib.urlencode({
                        'username': username,
                        'password': password}))
        else:
            session_crl.setopt(pycurl.URL, self.webapi_url)
        session_crl.perform()
        session_crl.close()

    def format_query(self, pubchem_query):
        '''pubchem_query is the pubchem style query dict as used in pubchem.py'''
        query_dic = {}
        if pubchem_query.has_key('mode') and pubchem_query['mode'] in ('1', '2', '3', '4'):
            query_dic['mode'] = pubchem_query['mode']
        else:
            query_dic['mode'] = '1'
        if pubchem_query.has_key('num_results'):
            query_dic['num_results'] = pubchem_query['num_results']
        if query_dic['mode'] != '4' and pubchem_query.has_key('query'):
            query_dic['smiles'] = pubchem_query['query']
        if pubchem_query.has_key('db'):
            query_dic['db'] = pubchem_query['db']
        else:
            query_dic['db'] = '1'
        if pubchem_query.has_key('substructure'):
            query_dic['mode'] = '2'
            query_dic['smiles'] = pubchem_query['substructure']
        if pubchem_query.has_key('similarity'):
            query_dic['mode'] = '3'
            query_dic['min_simi'] = pubchem_query['similarity']

        for k, v in pubchem_query.items():
            if not query_dic.has_key(k):
                query_dic[k] = v

        return query_dic

    def search(self, pubchem_query, download = False):

        query_dic = {
                'query_format': 'json',
                'json': json.dumps(self.format_query(pubchem_query))
                }

        self.fp.seek(0)
        self.fp.truncate()
        result_dic = {}

        self.crl_post.setopt(self.crl_post.POSTFIELDS,
                urllib.urlencode(query_dic).encode('utf8'))
        self.crl_post.perform()

        js = self.fp.getvalue()
        try:
            post_dic = json.loads(js)
        except:
            post_dic = {}

        results_dic = {}
        if post_dic.has_key('status'):
            if post_dic['status'] == 'OK' and post_dic.has_key('result_url'):
                self.fp.seek(0)
                self.fp.truncate()
                self.crl_get.setopt(pycurl.URL, str(self.host + post_dic['result_url']))
                self.crl_get.perform()
                try:
                    results_dic = json.loads(self.fp.getvalue())
                except:
                    return {}
                # filter isotopic molecules
                if results_dic.has_key('result_list'):
                    tmpl = []
                    for m in results_dic['result_list']:
                        if m.has_key('mol_2d_struc') and ( not 'ISO' in m['mol_2d_struc'] ):
                            tmpl.append(m)
                    results_dic['len_results'] = len(tmpl)
                    results_dic['results_list'] = tmpl
        if download:
            return results_dic
        return True

if __name__ == '__main__':
    import time

    q = Query(webapi_url)

    query = {
            'mode': '1',
            'db': '1',
            'mol': u'''\r\nChemJSDraw12161010372D\r\n\r\n  6  6  0  0  0  0  0  0  0  0  1 V2000\r\n   20.5140   -9.4120    0.0000 C   0  0\r\n   19.1630   -8.6320    0.0000 C   0  0\r\n   19.1630   -7.0720    0.0000 C   0  0\r\n   21.8650   -8.6320    0.0000 C   0  0\r\n   21.8650   -7.0720    0.0000 C   0  0\r\n   20.5140   -6.2920    0.0000 C   0  0\r\n  1  2  1  0\r\n  2  3  2  0\r\n  1  4  2  0\r\n  4  5  1  0\r\n  5  6  2  0\r\n  6  3  1  0\r\nM  END\r\n"'''
            }

    query = {
            'mode': '4',
            'db': '1',
            #'query': '''ems >= 310.0730 AND ems <= 310.0736 MAX = 20''',
            'query': 'mid = 1234 max = 1'
            }

    print q.search(query, True) #['logged_user']
    raw_input()
